import torch
import torch.nn as nn
import numpy as np
from SkyNet import *

net = torch.load('./FQSkyNet.pth')
net.cpu()
conv_layers = []
bn_layers = []
for module in net.named_modules():
    if isinstance(module[1], nn.Conv2d):
        conv_layers.append(module[1])
    if isinstance(module[1], nn.BatchNorm2d):
        bn_layers.append(module[1])

conv_layers[0].in_channels = 32
conv_layers[0].out_channels = 32
conv_layers[0].groups = 32
bn_layers[0].num_features = 32

conv_layers[1].in_channels = 32
conv_layers[1].out_channels = 64
bn_layers[1].num_features = 64

conv_layers[2].in_channels = 64
conv_layers[2].out_channels = 64
conv_layers[2].groups = 64
bn_layers[2].num_features = 64

conv_layers[3].in_channels = 64

weight = torch.cat([conv_layers[0].weight,torch.zeros(29,1,3,3)],0)
conv_layers[0].weight = nn.Parameter(weight)
conv_layers[0].bias = nn.Parameter(torch.cat([conv_layers[0].bias,torch.zeros(29)],0))
bn_layers[0].weight = nn.Parameter(torch.cat([bn_layers[0].weight,torch.zeros(29)],0))

weight = torch.cat([conv_layers[1].weight,torch.zeros(48,29,1,1)],1)
weight = torch.cat([weight,torch.zeros(16,32,1,1)],0)
conv_layers[1].weight = nn.Parameter(weight)
conv_layers[1].bias = nn.Parameter(torch.cat([conv_layers[1].bias,torch.zeros(16)],0))
bn_layers[1].weight = nn.Parameter(torch.cat([bn_layers[1].weight,torch.zeros(16)],0))

weight = torch.cat([conv_layers[2].weight,torch.zeros(16,1,3,3)],0)
conv_layers[2].weight = nn.Parameter(weight)
conv_layers[2].bias = nn.Parameter(torch.cat([conv_layers[2].bias,torch.zeros(16)],0))
bn_layers[2].weight = nn.Parameter(torch.cat([bn_layers[2].weight,torch.zeros(16)],0))

weight = torch.cat([conv_layers[3].weight,torch.zeros(96,16,1,1)],1)
conv_layers[3].weight = nn.Parameter(weight)

conv_layers[12].out_channels = 32
bn_layers[12].num_features = 32

weight = torch.cat([conv_layers[12].weight,torch.zeros(22,96,1,1)],0)
conv_layers[12].weight = nn.Parameter(weight)
conv_layers[12].bias = nn.Parameter(torch.cat([conv_layers[12].bias,torch.zeros(22)],0))
bn_layers[12].weight = nn.Parameter(torch.cat([bn_layers[12].weight,torch.zeros(22)],0))

print(net)
torch.save(net,"EQSkyNet.pth")